#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import pwd
import stat
import time
import signal
import logging
import psycopg
import lxml.etree
import subprocess
from mc_util import McUtil


class Storage:

    @staticmethod
    def get_properties():
        return {
            "need-temp-directory": True,
            "with-integrated-advertiser": True,
        }

    def __init__(self, param):
        self._listenIp = param["listen-ip"]
        self._tmpDir = param["temp-directory"]
        self._logDir = param["log-directory"]
        self._bAdvertiseDict = dict()                   # {mirror-site-id:bAdvertise}
        self._serverDict = dict()                       # {mirror-site-id:postgresql-server-object}

    def dispose(self):
        for msObj in self._serverDict.values():
            msObj.dispose()
        self._serverDict = dict()

    def add_mirror_site(self, mirror_site_id, mirror_site_properties):
        try:
            xmlElem = lxml.etree.fromstring(mirror_site_properties["config-xml"])

            # get advertise flag
            self._bAdvertiseDict[mirror_site_id] = (len(xmlElem.xpath(".//advertise")) > 0)

            # get template1 attributes: locale, encoding, extensions
            dbLocale = "C"
            dbEncoding = "UTF-8"
            dbExtensionList = []
            dbReadUser = mirror_site_id
            if len(xmlElem.xpath(".//template1")) > 0:
                xmlElem2 = xmlElem.xpath(".//template1")[0]
                if len(xmlElem2.xpath(".//locale")) > 0:
                    dbLocale = xmlElem2.xpath(".//locale")[0].text
                if len(xmlElem2.xpath(".//encoding")) > 0:
                    dbEncoding = xmlElem2.xpath(".//encoding")[0].text
                for extElem in xmlElem2.xpath(".//extension"):
                    dbExtensionList.append(extElem.text)
            if len(xmlElem.xpath(".//read-user")) > 0:
                dbReadUser = xmlElem.xpath(".//read-user")[0].text

            # create server objects
            self._serverDict[mirror_site_id] = _PostgresqlServer(self._listenIp, self._tmpDir, self._logDir,
                                                                 mirror_site_properties["data-directory"],
                                                                 dbLocale, dbEncoding, dbExtensionList, dbReadUser)

            # show log when the first database is being advertised
            if len(self._bAdvertiseDict) == 1:
                # here we can not give out port information
                logging.info("Advertiser (postgresql) started.")
        except Exception:
            if mirror_site_id in self._serverDict:
                self._serverDict[mirror_site_id].dispose()
                del self._serverDict[mirror_site_id]
            if mirror_site_id in self._bAdvertiseDict:
                del self._bAdvertiseDict[mirror_site_id]
            raise

    def advertise_mirror_site(self, mirror_site_id):
        self._serverDict[mirror_site_id].export()

    def remove_mirror_site(self, mirror_site_id):
        self._serverDict[mirror_site_id].dispose()
        del self._serverDict[mirror_site_id]
        del self._bAdvertiseDict[mirror_site_id]

    def get_param_for_updater(self, mirror_site_id):
        return {
            "unix-socket-file": self._serverDict[mirror_site_id].dbSocketFile,
            "files-directory": self._serverDict[mirror_site_id].dbFilesDir,
        }

    def get_param_for_advertiser(self, mirror_site_id):
        return {
            "unix-socket-file": self._serverDict[mirror_site_id].dbSocketFile,
        }

    def get_access_info(self, mirror_site_id):
        return {
            "url": "postgresql://%s@{IP}:%d/%s" % (self._serverDict[mirror_site_id].dbReadUser, self._serverDict[mirror_site_id].dbPort, mirror_site_id),
            "description": "",
        }


class _PostgresqlServer:

    def __init__(self, listenIp, tmpDir, logDir, dataDir, dbLocale, dbEncoding, dbExtensionList, dbReadUser):
        self._dbSuperUser = "postgres"
        self._dbReadUser = dbReadUser

        self._dbDataDir = dataDir
        self._dbInitFailedFile = os.path.join(self._dbDataDir, "initialize.failed")     # we create a temporary initialize.failed file
        self._dbFilesDir = os.path.join(self._dbDataDir, "_files")                      # we create a custom files directory in database's data directory

        self._tmpDir = tmpDir

        self._identFile = None
        self._hbaFile = None
        self._port = None
        self._proc = None
        self._pidFile = None
        self._optsFile = None
        self._autoConfFile = None
        self._confFile = None
        self._socketFile = None
        try:
            # initialize if needed
            if not self._isInitialized():
                self._initialize(os.path.join(logDir, "initdb.log"), os.path.join(logDir, "server.log"), dbLocale, dbEncoding, dbExtensionList)

            # prepare pg_ident.conf and pg_hba.conf
            self._identFile = os.path.join(self._dbDataDir, "pg_ident.conf")
            self._hbaFile = os.path.join(self._dbDataDir, "pg_hba.conf")
            self._writeIdentFileAndHbaFile(self._identFile, self._hbaFile, False)

            # prepare configuration files
            self._confFile = os.path.join(self._dbDataDir, "postgresql.conf")
            self._autoConfFile = os.path.join(self._dbDataDir, "postgresql.auto.conf")
            self._writeConfFile(self._confFile, self._autoConfFile)

            # allocate listening port
            self._port = McUtil.getFreeSocketPort("tcp")

            # start postgresql server process
            cmdList = [
                "postgres",
                "--listen_addresses=%s" % (listenIp),
                "--unix_socket_directories=%s" % (self._tmpDir),
                "-p", str(self._port),
                "-D", self._dbDataDir,
            ]
            self._proc = subprocess.Popen(cmdList, cwd=self._tmpDir)

            # get runtime files
            McUtil.waitSocketPortForProc("tcp", listenIp, self._port, self._proc)
            self._optsFile = McUtil.getExistFile(os.path.join(self._dbDataDir, "postmaster.opts"))
            self._socketFile = McUtil.findSocketFileInDir(self._tmpDir)
            self._pidFile = McUtil.getExistFile(os.path.join(self._dbDataDir, "postmaster.pid"))
        except BaseException:
            self.dispose()
            raise

    def dispose(self):
        if self._proc is not None:
            self._proc.terminate()
            self._proc.wait()
            self._proc = None
        if self._pidFile is not None:
            while os.path.exists(self._pidFile):
                time.sleep(0.1)
            self._pidFile = None
        if self._socketFile is not None:
            assert not os.path.exists(self._socketFile)
            self._socketFile = None
        if self._optsFile is not None:
            McUtil.forceDelete(self._optsFile)
            self._optsFile = None
        if self._confFile is not None:
            McUtil.forceDelete(self._confFile)
            self._confFile = None
        if self._autoConfFile is not None:
            McUtil.forceDelete(self._autoConfFile)
            self._autoConfFile = None
        if self._port is not None:
            self._port = None
        if self._hbaFile is not None:
            McUtil.forceDelete(self._hbaFile)
            self._hbaFile = None
        if self._identFile is not None:
            McUtil.forceDelete(self._identFile)
            self._identFile = None

    @property
    def dbSocketFile(self):
        return self._socketFile

    @property
    def dbPort(self):
        return self._port

    @property
    def dbReadUser(self):
        return self._dbReadUser

    @property
    def dbFilesDir(self):
        return self._dbFilesDir

    def export(self):
        assert self._identFile is not None and self._hbaFile is not None
        assert self._proc is not None

        self._writeIdentFileAndHbaFile(self._identFile, self._hbaFile, True)
        self._proc.send_signal(signal.SIGHUP)

    def _isInitialized(self):
        if os.path.exists(self._dbInitFailedFile):
            return False
        if not os.path.exists(os.path.join(self._dbDataDir, "PG_VERSION")):
            return False
        if not os.path.exists(self._dbFilesDir):
            return False
        return True

    def _initialize(self, initdbLogFile, serverLogFile, dbLocale, dbEncoding, dbExtensionList):
        identFile = os.path.join(self._dbDataDir, "pg_ident.conf")
        hbaFile = os.path.join(self._dbDataDir, "pg_hba.conf")
        autoConfFile = os.path.join(self._dbDataDir, "postgresql.auto.conf")
        confFile = os.path.join(self._dbDataDir, "postgresql.conf")
        optsFile = os.path.join(self._dbDataDir, "postmaster.opts")
        pidFile = os.path.join(self._dbDataDir, "postmaster.pid")

        McUtil.mkDirAndClear(self._dbDataDir)

        # create an empty database, set superuser's name to postgres
        try:
            out = McUtil.cmdCall("initdb", "--username=%s" % (self._dbSuperUser), "--locale=%s" % (dbLocale), "--encoding=%s" % (dbEncoding), self._dbDataDir)
            McUtil.touchFile(self._dbInitFailedFile)        # we should create initialize.failed before all actions, but have to do it here because initdb only accepts empty self._dbDataDir
            with open(initdbLogFile, "w") as f:
                f.write(out)
        except BaseException:
            McUtil.touchFile(self._dbInitFailedFile)        # same as above
            raise

        # create filesDir
        McUtil.mkDirAndClear(self._dbFilesDir)

        # create read user and write user
        proc = None
        try:
            # only allow os-user "mirrors" login to database server as user "postgres", without password
            with open(identFile, "w") as f:
                f.write("# MAPNAME  SYSTEM-USERNAME  PG-USERNAME\n")
                f.write("mirrors    mirrors          %s\n" % (self._dbSuperUser))
            with open(hbaFile, "w") as f:
                f.write("# TYPE  DATABASE  USER  ADDRESS  METHOD\n")
                f.write("local   all       all            ident map=mirrors\n")

            # run database server temporirily
            cmdStr = "postgres"
            cmdStr += " --unix_socket_directories='%s'" % (self._tmpDir)
            cmdStr += " --listen_addresses="
            cmdStr += " -D '%s'" % (self._dbDataDir)
            cmdStr += " >'%s' 2>&1" % (serverLogFile)
            proc = subprocess.Popen(cmdStr, cwd=self._tmpDir, shell=True)

            # wait database server start complete
            while True:
                bFound = False
                for fn in os.listdir(self._tmpDir):
                    if stat.S_ISSOCK(os.stat(os.path.join(self._tmpDir, fn)).st_mode):
                        bFound = True
                        break
                if bFound:
                    break
                time.sleep(0.1)

            # login to database and execute actions
            with psycopg.connect(user=self._dbSuperUser, host=self._tmpDir, dbname="template1") as conn:
                with conn.cursor() as c:
                    # write user
                    wname = pwd.getpwuid(os.getuid()).pw_name
                    c.execute("CREATE ROLE %s LOGIN CREATEDB;" % (wname))
                    c.execute("GRANT pg_write_all_data TO %s;" % (wname))               # FIXME: should not write system database, such as template0, template1, postgres

                    # read user
                    c.execute("CREATE ROLE %s LOGIN;" % (self._dbReadUser))
                    c.execute("GRANT pg_read_all_data TO %s;" % (self._dbReadUser))     # FIXME: should not read system database

                    # create extensions on database template1
                    if len(dbExtensionList) > 0:
                        for e in dbExtensionList:
                            c.execute("CREATE EXTENSION %s;" % (e))

        finally:
            if proc is not None:
                proc.terminate()
                proc.wait()
                proc = None
            if pidFile is not None:
                while os.path.exists(pidFile):
                    time.sleep(0.1)
                pidFile = None
            McUtil.forceDelete(confFile)
            McUtil.forceDelete(autoConfFile)
            McUtil.forceDelete(hbaFile)
            McUtil.forceDelete(identFile)
            McUtil.forceDelete(optsFile)

        os.unlink(self._dbInitFailedFile)

    def _writeIdentFileAndHbaFile(self, identFile, hbaFile, allowReadUser):
        wname = pwd.getpwuid(os.getuid()).pw_name

        with open(identFile, "w") as f:
            f.write("")

        with open(hbaFile, "w") as f:
            f.write("# TYPE  DATABASE        USER            ADDRESS                 METHOD\n")
            f.write("local   all             %s                                      trust\n" % (self._dbReadUser))
            f.write("local   all             %s                                      ident\n" % (wname))
            if allowReadUser:
                f.write("host    all             %s              0.0.0.0/0               trust\n" % (self._dbReadUser))

    def _writeConfFile(self, confFile, autoConfFile):
        # 1. use serializable isolation level globally
        with open(confFile, "w") as f:
            f.write("default_transaction_isolation = serializable\n")

        # we don't accept ALTER SYSTEM, this file should always be empty
        with open(autoConfFile, "w") as f:
            f.write("")
